import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import seaborn as sns

# 假设 agent_positions 是一个 (A, T, 2) 的数组
# A: 车辆数量, T: 时间步数, 2: 车辆的 (x, y) 坐标
A = 100  # 车辆数量
T = 50   # 时间步数
np.random.seed(0)

# 随机生成车辆的位置数据
agent_positions = np.random.rand(A, T, 2) * 100  # 车辆的位置在 0 到 100 的区域内

# 选择多个时间步 (例如 t=10 到 t=30)，合并这些时间步的数据
start_time = 10
end_time = 30
positions_at_multiple_times = agent_positions[:, start_time:end_time, :].reshape(-1, 2)

# 提取所有时间步中车辆的 (x, y) 坐标
x_multiple = positions_at_multiple_times[:, 0]
y_multiple = positions_at_multiple_times[:, 1]

# 定义网格大小，这里设置为 10x10 的网格
grid_size = 10  # 可以根据需要调整

# 计算热力图的密度
heatmap_multi, xedges_multi, yedges_multi = np.histogram2d(x_multiple, y_multiple, bins=grid_size, range=[[0, 100], [0, 100]])

# 创建网格坐标
xpos, ypos = np.meshgrid(xedges_multi[:-1] + (xedges_multi[1] - xedges_multi[0]) / 2,
                         yedges_multi[:-1] + (yedges_multi[1] - yedges_multi[0]) / 2)

xpos = xpos.flatten()
ypos = ypos.flatten()
zpos = np.zeros_like(xpos)  # 热力图的底部（z=0）

# 设置热力图的值
zsize = heatmap_multi.flatten()

# 绘制三维热力图
fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111, projection='3d')

# 绘制三维柱状图
ax.bar3d(xpos, ypos, zpos, dx=np.ones_like(zsize) * (xedges_multi[1] - xedges_multi[0]),
         dy=np.ones_like(zsize) * (yedges_multi[1] - yedges_multi[0]), dz=zsize, cmap='viridis')

# 设置标题和标签
ax.set_title(f'3D Heatmap of Vehicle Distribution from Time {start_time} to {end_time}')
ax.set_xlabel('X Position')
ax.set_ylabel('Y Position')
ax.set_zlabel('Vehicle Density')

plt.show()
